{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "🦙 Chat Templates in HuggingFace is just great. 🔥🚀\n",
    "\n",
    "📌 Why templates? ❓\n",
    "\n",
    "An increasingly common use case for LLMs is chat. In a chat context, rather than continuing a single string of text (as is the case with a standard language model), the model instead continues a conversation that consists of one or more messages, each of which includes a role as well as message text.\n",
    "\n",
    "Chat models have been trained with very different formats for converting conversations into a single tokenizable string. Using a format different from the format a model was trained with will usually cause severe, silent performance degradation, so matching the format used during training is extremely important! \n",
    "\n",
    "So, `apply_chat_template` attribute that can be used to save the chat format the model was trained with. This attribute contains a Jinja template that converts conversation histories into a correctly formatted string. \n",
    "\n",
    "--------\n",
    "\n",
    "All language models, including models fine-tuned for chat, operate on linear sequences of tokens and do not intrinsically have special handling for roles. This means that role information is usually injected by adding control tokens between messages, to indicate both the message boundary and the relevant roles.\n",
    "\n",
    "Unfortunately, there isn’t (yet!) a standard for which tokens to use, and so different models have been trained with wildly different formatting and control tokens for chat. This can be a real problem for users - if you use the wrong format, then the model will be confused by your input, and your performance will be a lot worse than it should be. \n",
    "\n",
    "Whether you're fine-tuning a model or using it directly for inference, it's always a good idea to minimize these distribution shifts and keep the input you give it as similar as possible to the input it was trained on. With regular language models, it's relatively easy to do that - simply load your tokenizer and model from the same checkpoint, and you're good to go.\n",
    "\n",
    "With chat models, however, it's a bit different. This is because \"chat\" is not just a single string of text that can be straightforwardly tokenized - it's a sequence of messages, each of which contains a role as well as content, which is the actual text of the message. Most commonly, the roles are \"user\" for messages sent by the user, \"assistant\" for responses written by the model, and optionally \"system\" for high-level directives given at the start of the conversation.\n",
    "\n",
    "This sequence of messages needs to be converted into a text string before it can be tokenized and used as input to a model. \n",
    "\n",
    "This is the problem that chat templates aim to solve. Chat templates are Jinja template strings that are saved and loaded with your tokenizer, and that contain all the information needed to turn a list of chat messages into a correctly formatted input for your model. Here are three chat template strings, corresponding to the three message formats above:\n",
    "\n",
    "------------------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Chat Templates in HuggingFace is just great 🚀\n",
    "\n",
    "from transformers import AutoTokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"meta-llama/Llama-2-7b-chat-hf\")\n",
    "\n",
    "chat = [\n",
    "  {\"role\": \"user\", \"content\": \"Hello, how are you?\"},\n",
    "  {\"role\": \"assistant\", \"content\": \"I'm doing great. How can I help you today?\"},\n",
    "  {\"role\": \"user\", \"content\": \"I'd like to show off how chat templating works!\"},\n",
    "]\n",
    "\n",
    "tokenizer.use_default_system_prompt = False\n",
    "tokenizer.apply_chat_template(chat, tokenize=False)\n",
    "\n",
    "\"<s>[INST] Hello, how are you? [/INST] I'm doing great. How can I help you today? </s><s>[INST] \\\n",
    "    I'd like to show off how chat templating works! [/INST]\"\n",
    "\n",
    "# Note that the tokenizer has added the control tokens [INST] and [/INST]\n",
    "# to indicate the start and end of user messages\n",
    "\n",
    "\n",
    "# Now with Jinja template to format inputs similarly to the way LLaMA formats them\n",
    "\n",
    "{% for message in messages %}\n",
    "    {% if message['role'] == 'user' %}\n",
    "        {{ bos_token + '[INST] ' + message['content'].strip() + ' [/INST]' }}\n",
    "    {% elif message['role'] == 'system' %}\n",
    "        {{ '<<SYS>>\\\\n' + message['content'].strip() + '\\\\n<</SYS>>\\\\n\\\\n' }}\n",
    "    {% elif message['role'] == 'assistant' %}\n",
    "        {{ '[ASST] '  + message['content'] + ' [/ASST]' + eos_token }}\n",
    "    {% endif %}\n",
    "{% endfor %}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "📌 How do I get started with templates?\n",
    "\n",
    "Easy! If a tokenizer has the chat_template attribute set, it's ready to go. You can use that model and tokenizer in `ConversationPipeline`, or you can call `tokenizer.apply_chat_template()` to format chats for inference or training. \n",
    "\n",
    "------------------\n",
    "\n",
    "📌  How do I create a new chat template?\n",
    "\n",
    "You can add a `chat_template` even for checkpoints that you're not the owner of, by opening a pull request. The only change you need to make is to set the tokenizer.`chat_template` attribute to a Jinja template string. Once that's done, push your changes and you're ready to go!\n",
    "\n",
    "--------------\n",
    "\n",
    "📌 How do chat templates work?\n",
    "\n",
    "The chat template for a model is stored on the `tokenizer.chat_template` attribute. If no chat template is set, the default template for that model class is used instead.\n",
    "\n",
    "Jinja template gives you a lot of flexibility - Let’s see a Jinja template that can format inputs similarly to the way LLaMA formats them \n",
    "\n",
    "```\n",
    "{% for message in messages %}\n",
    "    {% if message['role'] == 'user' %}\n",
    "        {{ bos_token + '[INST] ' + message['content'].strip() + ' [/INST]' }}\n",
    "    {% elif message['role'] == 'system' %}\n",
    "        {{ '<<SYS>>\\\\n' + message['content'].strip() + '\\\\n<</SYS>>\\\\n\\\\n' }}\n",
    "    {% elif message['role'] == 'assistant' %}\n",
    "        {{ '[ASST] '  + message['content'] + ' [/ASST]' + eos_token }}\n",
    "    {% endif %}\n",
    "{% endfor %}\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## From HF Doc\n",
    "\n",
    "```py\n",
    "from transformers import AutoTokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"mistralai/Mistral-7B-Instruct-v0.1\")\n",
    "\n",
    "chat = [\n",
    "  {\"role\": \"user\", \"content\": \"Hello, how are you?\"},\n",
    "  {\"role\": \"assistant\", \"content\": \"I'm doing great. How can I help you today?\"},\n",
    "  {\"role\": \"user\", \"content\": \"I'd like to show off how chat templating works!\"},\n",
    "]\n",
    "\n",
    "tokenizer.apply_chat_template(chat, tokenize=False)\n",
    "\n",
    "# \"<s>[INST] Hello, how are you? [/INST]I'm doing great. How can I help you today?</s> [INST] I'd like to show off how chat templating works! [/INST]\"\n",
    "\n",
    "```\n",
    "\n",
    "\n",
    "Notice how the entire chat is condensed into a single string. If we use `tokenize=True`, which is the default setting, that string will also be tokenized for us. \n",
    "\n",
    "Note that for Mistral-7B-Instruct, the tokenizer has added the control tokens [INST] and [/INST] to indicate the start and end of user messages (but not assistant messages!). Mistral-instruct was trained with these tokens, but BlenderBot was not."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## More on why we need `apply_chat_template`\n",
    "\n",
    "Basically, it auto-applies the right formatting around the messages for models.\n",
    "\n",
    "Using this could greatly improve the user experience of open-source LLMs since users would no longer have to manually add the correct tokens and formatting to prompts sent to the PromptNode.\n",
    "\n",
    "Here is a full example of how to use this new function with the Mistral Instruct Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "device = \"cuda\" # the device to load the model onto\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\"mistralai/Mistral-7B-Instruct-v0.1\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"mistralai/Mistral-7B-Instruct-v0.1\")\n",
    "\n",
    "messages = [\n",
    "    {\"role\": \"user\", \"content\": \"What is your favourite condiment?\"},\n",
    "    {\"role\": \"assistant\", \"content\": \"Well, I'm quite partial to a good squeeze of fresh lemon juice. It adds just the right amount of zesty flavour to whatever I'm cooking up in the kitchen!\"},\n",
    "    {\"role\": \"user\", \"content\": \"Do you have mayonnaise recipes?\"}\n",
    "]\n",
    "\n",
    "encodeds = tokenizer.apply_chat_template(messages, return_tensors=\"pt\")\n",
    "\n",
    "model_inputs = encodeds.to(device)\n",
    "model.to(device)\n",
    "\n",
    "generated_ids = model.generate(model_inputs, max_new_tokens=1000, do_sample=True)\n",
    "decoded = tokenizer.batch_decode(generated_ids)\n",
    "print(decoded[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "------------------\n",
    "\n",
    "## Below is a common and useful util method for formatting the dialogue history into a prompt of a chat-LLM model. 🐍\n",
    "\n",
    "This will help understand the context of the conversation and generate appropriate responses by the Model\n",
    "\n",
    "--------\n",
    "\n",
    "The function takes a history of dialogues as input, which is a list of lists where each sublist represents a pair of user and assistant messages, like below\n",
    "\n",
    "```\n",
    "[\n",
    "    [\"User's first instruction\", \"Assistant's first response\"],\n",
    "    [\"User's second instruction\", \"Assistant's second response\"],\n",
    "    [\"User's third instruction\", null]\n",
    "]\n",
    "\n",
    "```\n",
    "\n",
    "And, the messages list should be of the following format:\n",
    "\n",
    "```\n",
    "messages =\n",
    "\n",
    "[\n",
    "    {\"role\": \"user\", \"content\": \"User's first message\"},\n",
    "    {\"role\": \"assistant\", \"content\": \"Assistant's first response\"},\n",
    "    {\"role\": \"user\", \"content\": \"User's second message\"},\n",
    "    {\"role\": \"assistant\", \"content\": \"Assistant's second response\"},\n",
    "    {\"role\": \"user\", \"content\": \"User's third message\"}\n",
    "]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_chat_history(history) -> str:\n",
    "    messages = [{\"role\": (\"user\" if i % 2 == 0 else \"assistant\"), \"content\": dialog[i % 2]}\n",
    "                for i, dialog in enumerate(history) for _ in (0, 1) if dialog[i % 2]]\n",
    "    # The conditional `(if dialog[i % 2])` ensures that messages\n",
    "    # that are None (like the latest assistant response in an ongoing\n",
    "    # conversation) are not included.\n",
    "    return pipeline.tokenizer.apply_chat_template(\n",
    "        messages, tokenize=False,\n",
    "        add_generation_prompt=True)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
